package com.stress.testing.plugin.java.jdbc.sql.util;

import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.stress.testing.common.StressTesting;
import com.stress.testing.common.TraceContextHolder;
import com.stress.testing.common.conf.AgentConfig;
import com.stress.testing.common.conf.PluginConfig;
import com.stress.testing.plugin.java.jdbc.sql.parser.TableNameVisitor;

import java.io.StringWriter;
import java.lang.reflect.Constructor;

import static com.stress.testing.common.conf.PluginConfigKey.DEFAULT_JAVA_JDBC_TABLE_NAME_PREFIX;
import static com.stress.testing.common.conf.PluginConfigKey.PLUGIN_JAVA_JDBC_SQL_PARSER_PREFIX;
import static com.stress.testing.common.conf.PluginConfigKey.PLUGIN_JAVA_JDBC_TABLE_NAME_PREFIX_KEY;

public final class SQLParserUtils {
    private final static String SQL_PARSER_UNKNOW = "UNKNOW";

    private SQLParserUtils() {
    }

    public static String changeTableName(String sql, AgentConfig agentConfig,
                                         PluginConfig pluginConfig, String databaseType) {
        if (!StressTesting.enable(agentConfig, pluginConfig, TraceContextHolder.getTraceContext())) {
            return sql;
        }

        String sqlParser = pluginConfig.get(PLUGIN_JAVA_JDBC_SQL_PARSER_PREFIX + databaseType, SQL_PARSER_UNKNOW);
        if (SQL_PARSER_UNKNOW.equals(sqlParser)) {
            return sql;
        }

        SQLStatementParser parser = null;
        try {
            Class<?> sqlParserClass = SQLParserUtils.class.getClassLoader().loadClass(sqlParser);
            Constructor<?> constructor = sqlParserClass.getConstructor(String.class);
            Object instance = constructor.newInstance(sql);
            parser = SQLStatementParser.class.cast(instance);
        } catch (Exception e) {
            // warning log
            return sql;
        }

        if (parser == null) {
            return sql;
        }

        // TODO: 可以通过本地缓存优化
        StringWriter appender = new StringWriter();
        String tableNamePrefix = pluginConfig.get(PLUGIN_JAVA_JDBC_TABLE_NAME_PREFIX_KEY,
                DEFAULT_JAVA_JDBC_TABLE_NAME_PREFIX);
        TableNameVisitor visitor = new TableNameVisitor(appender, tableNamePrefix);

        SQLStatement sqlStatement = parser.parseStatement();
        sqlStatement.accept(visitor);

        return appender.toString();
    }
}
